/*
 * main.cpp
 *
 *  Created on: Dec 11, 2013
 *      Author: thanhkm
 */

#include "Features.h"
#include "IO.h"
#include "Learning.h"
#include "BOWScene.h"

#define DIC_FILE_PRE "trained_data/dictionary"
#define DIC_NAME "dictionary"
#define HISTOS_FILE_PRE "trained_data/histograms"
#define HISTOS_NAME "histograms"
#define HISTOS_TEST_FILE_PRE "trained_data/histograms_test"
#define HISTOS_TEST_NAME "histograms_test"
#define LABELS_FILE_PRE "trained_data/labels"
#define LABELS_NAME "labels"
#define LABELS_TEST_FILE_PRE "trained_data/labels_test"
#define LABELS_TEST_NAME "labels_test"
#define DICTIONARY_SIZE 200
#define DICTIONARY_BUILD 1

int main(int argc, char* argv[]) {
	std::string train_csv;
	std::string test_csv;
	std::string dic_file;
	std::string histos_file;
	std::string labels_file;
	std::string histos_test_file;
	std::string labels_test_file;

	// read arguments
	train_csv = std::string(argv[1]);
	test_csv = std::string(argv[2]);
	int dictionarySize = atoi(argv[3]);
	dic_file = std::string(DIC_FILE_PRE) + std::string(argv[3]) +  std::string(".yml");
	histos_file = std::string(HISTOS_FILE_PRE) + std::string(argv[3]) +  std::string(".yml");
	histos_test_file = std::string(HISTOS_TEST_FILE_PRE) + std::string(argv[3]) +  std::string(".yml");
	labels_file = std::string(LABELS_FILE_PRE) + std::string(".yml");
	labels_test_file = std::string(LABELS_TEST_FILE_PRE) + std::string(".yml");
	int pha = atoi(argv[4]);

//	std::cout << dic_file << std::endl
//			  << histos_file << std::endl
//			  << histos_test_file << std::endl
//			  << labels_file << std::endl
//			  << labels_test_file << std::endl;

	std::vector<cv::Mat> trainImages, testImages;
	std::vector<int> trainLabels, testLabels;
	cv::Mat dictionary;
	cv::Mat mTrainHistograms, mTestHistograms;
	cv::Mat mTrainLabels, mTestLabels;

	// read train images
	IO::read_csv(train_csv, trainImages, trainLabels);
	// read test images
	IO::read_csv(test_csv, testImages, testLabels);

	BOWScene bow;
	Features fea;

if (pha == 0){
	//------------------------------------------------------------
	// pha 0: compute the dictionary (bag of words)
	//------------------------------------------------------------
	// compute descriptors of all images
	fea.compute(trainImages);
	std::cout << fea.numberKeypoints() << std::endl;
	std::vector<cv::Mat> vecDescriptors = fea.getDescritpors();
	// compute the bag of words
	// and histograms of all images
	bow.cluster(vecDescriptors, dictionarySize);
	dictionary = bow.getDictionary();
	// save dictionary
	IO::writeMat(dic_file.c_str(), dictionary, DIC_NAME);

} else if (pha == 1){
	//------------------------------------------------------------
	// pha 1: compute histogram of images following the dictionary
	//------------------------------------------------------------
	//read dictionary
	IO::readMat(dic_file.c_str(), dictionary, DIC_NAME);
	// compute histograms
	std::cout << "compute histogram" << std::endl;
	bow.computeHistogram(trainImages, mTrainHistograms, dictionary);
	// convert trainLabels to labels
	Learning::convertToMat(trainLabels, mTrainLabels);
	// save histograms of all images and theirs labels
	IO::writeMat(histos_file.c_str(), mTrainHistograms, HISTOS_NAME);
	IO::writeMat(labels_file.c_str(), mTrainLabels, LABELS_NAME);
} else if (pha == 2){
	//------------------------------------------------------------
	// pha 2:
	// compute histogram of test images following the dictionary
	// then save it in to file
	//------------------------------------------------------------
	//read dictionary
	IO::readMat(dic_file.c_str(), dictionary, DIC_NAME);
	// compute bag of words histograms
	bow.setDictionary(dictionary);
	bow.computeHistogram(testImages, mTestHistograms, dictionary);
	std::cout << mTestHistograms.size() << std::endl;
	IO::writeMat(histos_test_file.c_str(), mTestHistograms, HISTOS_TEST_NAME);
	// convert trainLabels to labels
	Learning::convertToMat(testLabels, mTestLabels);
	IO::writeMat(labels_test_file.c_str(), mTestLabels, LABELS_TEST_NAME);
} else if (pha == 3){
	//------------------------------------------------------------
	// pha 3:
	// training SVM with histograms calculated
	// then test the recognition
	//------------------------------------------------------------
	//read histogram
	IO::readMat(histos_file.c_str(), mTrainHistograms, HISTOS_NAME);
	//read labels
	IO::readMat(labels_file.c_str(), mTrainLabels, LABELS_NAME);
	//---------------------------------------------------------
	//read histogram test
	IO::readMat(histos_test_file.c_str(), mTestHistograms, HISTOS_TEST_NAME);
	//read labels test
	IO::readMat(labels_test_file.c_str(), mTestLabels, LABELS_TEST_NAME);
	//---------------------------------------------------------
	// training with the data
	Learning learn;
	learn.train(mTrainHistograms, mTrainLabels);

	// test the program
	cv::Mat predicted;
	learn.predict(mTestHistograms, predicted);

	// convert trainLabels to labels
	Learning::convertToMat(testLabels, mTestLabels);
	IO::writeMat(labels_test_file.c_str(), mTestLabels, LABELS_TEST_NAME);
	// print result
	std::cout << "confusion matrix" << std::endl;
	Learning::evaluate(predicted, mTestLabels);
}
	return 0;
}
